import sys
import random
import pylab as plt
import numpy as np
import math
import os
import networkx as nx

from params import *

datasets = ['bkite','fsquare','gowalla']
lbls = ['Brightkite', 'Foursquare', 'Gowalla']
node_dist = [5651.0,8494.0,5663.0]
link_dist = [2041.0,1442.0,1792.0]
degs = [7.88,22.07,9.48,7.0]

plots = []
for dataset in datasets:
    trace_file = os.path.join(WORKDIR, 'gsn', 'traces',dataset, '%s_graph.txt'%dataset)
    triangle_file = os.path.join(WORKDIR, 'gsn','results',dataset,'%s_triangles.txt'%dataset)
    print 'Dataset %s'%dataset

    BINSIZE = 10.0
    M1 = 0.0
    M2 = 4.0
    MIN_LENGTH = 0
    MAX_LENGTH = 10**M2
    STEPS = 100

    def scale_distance2(d):
        return BINSIZE*(1+int(float(d)//BINSIZE)) - BINSIZE/2

    def scale_distance(d):
        d = max(1.0,d)
        l = math.log10(d)
        ratio = (l-M1)/(M2-M1)
        index = float(int(ratio*STEPS))/STEPS
        f = M1 + index*(M2-M1)
        #print d,l, ratio,index,f
        return 10**f

    link_length = {}
    g = nx.Graph()
    for line in open(trace_file):
        if line.startswith('#'):
            continue
        u1,u2,dist = line.split(' ')
        dist = float(dist)
        u1,u2 = map(int,(u1,u2))
        g.add_edge(u1,u2,weight=dist)
        if MIN_LENGTH <= dist <= MAX_LENGTH:
            d = scale_distance(dist)
            link_length.setdefault(d,0)
            link_length[d] += 1

    print 'There are %d links'%g.size()
    link_triangles = set()
    tr = 0
    for line in open(triangle_file):
        n1,n2,n3 = sorted(map(int,line.split()))
        link_triangles.add((n1,n2))
        link_triangles.add((n1,n3))
        link_triangles.add((n2,n3))
        tr += 1
    print 'There are %d links in %d triangles'%(len(link_triangles),tr)
    print 'Average probability ', 1.0*len(link_triangles)/g.size()
    avg_tri_prob = 1.0*len(link_triangles)/g.size()

    triangle_length = {}
    for n1,n2 in link_triangles:
        dist = g[n1][n2]['weight']
        if MIN_LENGTH <= dist <= MAX_LENGTH:
            d = scale_distance(dist)
            triangle_length.setdefault(d,0)
            triangle_length[d] += 1

    x = sorted(link_length)
    links = [link_length[k] for k in x]
    triangles = []
    for k in x:
        if k in triangle_length:
            triangles.append(triangle_length[k])
        else:
            triangles.append(0.0)

    probs = []
    for a,b,c in zip(links,triangles,x):
        #print 'Length %f: triangles %d, links %d'%(c,b,a)
        probs.append(1.0*b/a)

    plots.append((x,probs,avg_tri_prob))

plt.figure()
plt.clf()
plt.axes(FIG_AXES2)
legs = []
for (x,p,avg_tri_prob),m in zip(plots,markers):
    ax1 = plt.semilogx(x,p, 'k%s'%m, mfc='None')
    plt.semilogx([x[0],x[-1]],[avg_tri_prob,avg_tri_prob], 'k-',
            linewidth=1)
    legs.append(ax1[0])
plt.legend(legs,lbls,loc='lower left',
        numpoints=1)
plt.xlabel('Link length [km]')
plt.ylabel('Probability')
plt.ylim(ymin=0.5, ymax=0.95)
#plt.legend(['Data','Average probability'],
#        loc='best',numpoints=1)
plt.grid(True)
#plt.savefig('%s_triangle_prob.pdf'%dataset)
plt.savefig('triangle_prob.pdf')
plt.close()
